from lib.base import BaseMetric
import torch
from torch.nn import functional as F
import numpy as np
from lib.utils import DBindex

class BaselineAcc(BaseMetric):
    def __init__(self, cfg):
        super().__init__(cfg)
        self.cfg = cfg



    def __call__(self, logits, labels, type):
        # logits_q: 3000*dim

        if type == 'TEST':
            return self.episodic_acc(logits, labels, type)

        if type == 'VALIDATE':
            class_file = {}
            labels = labels.cpu().numpy()
            logits = logits.detach().cpu().numpy()
            for logit, label in zip(logits, labels):
                if label not in class_file.keys():
                    class_file[label] = []
                class_file[label].append(logit)

            for cl in class_file:
                class_file[cl] = np.array(class_file[cl])
            a = DBindex(class_file)
            return 1 / a

        if type == 'TRAIN':
            return self.acc(logits, labels)
